#include "raytracer.h"
#include "material.h"
#include "vectors.h"
#include "argparser.h"
#include "raytree.h"
#include "utils.h"
#include "mesh.h"
#include "face.h"
#include "primitive.h"
#include "photon_mapping.h"

// ===========================================================================
// casts a single ray through the scene geometry and finds the closest hit
int RayTracer::CastRay(Ray &ray, Hit &h, bool use_rasterized_patches) const {
  int answer = 0;

  // intersect each of the quads
  for (int i = 0; i < mesh->numOriginalQuads(); i++) {
    Face *f = mesh->getOriginalQuad(i);
    if (f->intersect(ray,h,args->intersect_backfacing)) answer = true;
  }

  // intersect each of the primitives (either the patches, or the original primitives)
  if (use_rasterized_patches) {
    for (int i = 0; i < mesh->numRasterizedPrimitiveFaces(); i++) {
      Face *f = mesh->getRasterizedPrimitiveFace(i);
      if (f->intersect(ray,h,args->intersect_backfacing)) answer = true;
    }
  } else {
    int num_primitives = mesh->numPrimitives();
    for (int i = 0; i < num_primitives; i++) {
      int res = mesh->getPrimitive(i)->intersect(ray, h);
      if (res)
	answer = res;
    }
  }
  return answer;
}

Ray get_reflected_ray(Vec3f direction, Vec3f normal, Vec3f point)
{
  float c1 = -(normal.Dot3(direction));
  return Ray(point, direction+(2*normal*c1));
}

Ray get_refracted_ray(Vec3f direction, Vec3f normal, Vec3f point, double index1, double index2)
{
  float c1 = -(normal.Dot3(direction));
  float n = index1/index2;
  float c2 = sqrt(1 - n*n * (1 - c1*c1));
  Vec3f Rr = (n * direction) + ((n * c1 - c2) * normal);
  return Ray(point, Rr);
}

Vec3f RayTracer::TraceShadow(Ray &ray, Vec3f point, Face *f, Hit &hit) const
{
  Vec3f avg(0,0,0);
  Material *m = hit.getMaterial();
 
  for (int i = 0; i < args->num_shadow_samples; i++)
    {
      Hit temp = Hit();
      Vec3f pointOnLight = args->num_shadow_samples == 1 ? f->computeCentroid() : f->GridPoint(i/4,(i-i/4)%4);
      Vec3f lightColor = f->getMaterial()->getEmittedColor() * f->getArea();
      Vec3f dirToLight = pointOnLight - point;
      double dist = dirToLight.Length();
      dirToLight.Normalize();
      Ray light_ray = Ray(point, dirToLight);
      lightColor /= M_PI*dist*dist;
      
      RayTree::AddShadowSegment(light_ray, 0, dist);
      
      // ===========================================
      // ASSIGNMENT:  ADD SHADOW & SOFT SHADOW LOGIC
      // ===========================================
      // add the lighting contribution from this particular light at this point
      // (fix this to check for blockers between the light & this surface)
      //    if (!CastRay(light_ray, temp, false))
      //      {
      //    	answer += m->Shade(ray,hit,dirToLight,lightColor,args);
      //      }
      
      CastRay (light_ray, temp, false);
      if (fabs(temp.getT()-dist) <= 0.00001)
	avg += m->Shade(ray, hit, dirToLight, lightColor, args);
    }
  
  avg /= (args->num_shadow_samples);  
  return avg;
}

Vec3f RayTracer::TracePath(Ray &ray, Hit &hit, int bounce_count, double lindex) const {
  Material *m;
  int intersect;
  Vec3f emittance, reflected(0,0,0), diffuse, point, answer(0,0,0);
  Vec3f normal = hit.getNormal();
  float dp = 0.0;
  
  hit = Hit();
  intersect = CastRay(ray,hit,false);

  if (intersect == false){
		
	if (bounce_count >= 1){
	  return args->ambient_light;
	}
	if (bounce_count == 0){
	  return Vec3f(srgb_to_linear(args->background_color.r()),
		       srgb_to_linear(args->background_color.g()),
		       srgb_to_linear(args->background_color.b()));
	} 
	else {
	  return Vec3f(0,0,0);
	}
    
  }

  normal = hit.getNormal();
  point = ray.pointAtParameter(hit.getT()) + Vec3f(0.001,0.001,0.001)*normal;
  
  m = hit.getMaterial();
  assert (m != NULL);
  
  
  diffuse = m->getDiffuseColor(hit.get_s(),hit.get_t());
  emittance = m->getEmittedColor();
  
  if (emittance.Length() > 0.001) {
    return emittance;
  }

  double p = diffuse.x()>diffuse.y() && diffuse.x()>diffuse.z() ? diffuse.x() : diffuse.y()>diffuse.z() ? diffuse.y() : diffuse.z();

  if ((bounce_count > args->num_bounces && (RandomDouble() < p)) || (bounce_count > 2*args->num_bounces)) {
    return Vec3f(0,0,0);
  }
  diffuse = diffuse * (1/p);
  
  
  double davg = (diffuse.x()+diffuse.y()+diffuse.z())/3.0;

  if (m->getRefractionIndex() != 0)
    {
      double index = m->getRefractionIndex();
      double n = lindex / index;
      Vec3f N = normal * intersect;
      double cosI = -N.Dot3(ray.getDirection());
      double cosT2 = 1.0f-n*n*(1.0f-cosI*cosI);
      
      if (cosT2 > 0.0f)
	{
	  Vec3f dir = (n * ray.getDirection()) +
	    (n*cosI-sqrtf(cosT2))*N;
	  Ray n = Ray(point+dir*.00001,dir);
	  Hit t;
	  Vec3f rcol, absorbance, transparency;
	  
	  rcol = TracePath(n,t,bounce_count+1, index);

	  absorbance = diffuse * 0.15 * -t.getT();
	  transparency = Vec3f(expf(absorbance.x()),
			       expf(absorbance.y()),
			       expf(absorbance.z()));
	  answer += rcol * transparency;
	}
    }
  if (diffuse.Length() > 0)
    {
      Vec3f dir = RandomDiffuseDirectionHemisphere(hit.getNormal());
      Ray n = Ray(point+dir*.00001,dir);
      Hit t = Hit();
      

      reflected = TracePath(n, t, bounce_count+1, lindex);
      dp = n.getDirection().Dot3(hit.getNormal());
      
      answer += Vec3f(diffuse.x()*reflected.x()*dp,
			       diffuse.y()*reflected.y()*dp,
			       diffuse.z()*reflected.z()*dp);
    }
  if (m->getReflectiveColor().Length()>0)
    {
      Vec3f dir = MirrorDirection(hit.getNormal(),ray.getDirection()), BRDF;
      Ray n = Ray(point+dir*.00001,dir);
      Hit t;  
      
      BRDF = m->getReflectiveColor();
      
      reflected = Vec3f(0,0,0);
      
	  
      if ((m->getGloss() > 0) && (args->num_gloss_samples > 0) && (bounce_count < 1) ){
	double gloss = m->getGloss();
	
		Vec3f b1 = Vec3f(dir.z(),dir.y(),-dir.x());
		Vec3f b2;
		
		Vec3f::Cross3(b2,dir,b1);
		
		for (int i = 0; i < args->num_gloss_samples; i++){
		  Vec3f r, refres;
			float xoffs, yoffs;
			Hit temp;
			do {
				xoffs = RandomDouble() * gloss;
				yoffs = RandomDouble() * gloss;
			} while((xoffs*yoffs+yoffs*yoffs) > (gloss*gloss));
		  
			r = dir+b1*xoffs+b2*yoffs;
			r.Normalize();
		  
			Ray cast = Ray(point+r*.00001, r);
		  
			refres = TracePath(cast, temp, bounce_count+1, lindex);
			
			//			refres.Scale(BRDF.x(), BRDF.y(), BRDF.z());
			reflected += 1.0/args->num_gloss_samples * refres;

			
		}
      
		}
		else{
		  reflected = TracePath(n, t, bounce_count+1, lindex);

		}
      dp = 1.0;
      answer += Vec3f(reflected.x()*BRDF.x()*dp,dp*reflected.y()*BRDF.y(),reflected.z()*BRDF.z()*dp);
    }
  return answer;
}

// ===========================================================================
// does the recursive (shadow rays & recursive rays) work
Vec3f RayTracer::TraceRay(Ray &ray, Hit &hit, int bounce_count) const {

  // First cast a ray and see if we hit anything.
  hit = Hit();
  bool intersect = CastRay(ray,hit,false);
  
  // if there is no intersection, simply return the background color
  if (intersect == false) {
    return Vec3f(srgb_to_linear(args->background_color.r()),
		 srgb_to_linear(args->background_color.g()),
		 srgb_to_linear(args->background_color.b()));
    
  }

  // otherwise decide what to do based on the material
  Material *m = hit.getMaterial();
  assert (m != NULL);

  // rays coming from the light source are set to white, don't bother to ray trace further.
  if (m->getEmittedColor().Length() > 0.001) {
    return Vec3f(1,1,1);
  } 
 
  
  Vec3f normal = hit.getNormal();
  Vec3f point = ray.pointAtParameter(hit.getT()) + Vec3f(0.001,0.001,0.001)*normal;
  Vec3f answer;

  // ----------------------------------------------
  // ambient light
  Vec3f diffuse_color = m->getDiffuseColor(hit.get_s(),hit.get_t());
  if (args->gather_indirect) {
    // photon mapping for more accurate indirect light
    answer = diffuse_color * photon_mapping->GatherIndirect(point, normal, ray.getDirection());
  } else {
    // the usual ray tracing hack for indirect light
    answer = diffuse_color * args->ambient_light;
  }      

  // ----------------------------------------------
  // add contributions from each light that is not in shadow
  int num_lights = mesh->getLights().size();
  for (int i = 0; i < num_lights; i++) {
    answer += TraceShadow(ray, point, mesh->getLights()[i], hit);
  }
      
  // ----------------------------------------------
  // add contribution from reflection, if the surface is shiny
  Vec3f reflectiveColor = m->getReflectiveColor();

  if ((m->getGloss() > 0) && (bounce_count < 2))
    {
      double gloss = m->getGloss();
      Vec3f ref = ray.getDirection()-(2.0f*ray.getDirection().Dot3(normal))*normal;
      Vec3f b1 = Vec3f(ref.z(),ref.y(),-ref.x());
      Vec3f b2;
      
      Vec3f::Cross3(b2,ref,b1);
      
      for (int i = 0; i < args->num_gloss_samples; i++)
	{
	  Vec3f r, refres;
	  float xoffs, yoffs;
	  Hit temp;
	  do {
	    xoffs = RandomDouble() * gloss;
	    yoffs = RandomDouble() * gloss;
	  } while((xoffs*yoffs+yoffs*yoffs) > (gloss*gloss));
	  
	  r = ref+b1*xoffs+b2*yoffs;
	  r.Normalize();
	  
	  Ray cast = Ray(point+r*.000001, r);
	  
	  refres = TraceRay(cast, temp, bounce_count+1);
	  refres.Scale(reflectiveColor.x(), reflectiveColor.y(), reflectiveColor.z());
	  answer += 1.0/args->num_gloss_samples * refres;

	  RayTree::AddReflectedSegment(ray, 0, temp.getT());
	}
      
    }
  else if (bounce_count < args->num_bounces)
    {
      Hit temp;
      Ray reflected_ray = get_reflected_ray(ray.getDirection(), normal, point);
      Vec3f refres;
      refres = TraceRay(reflected_ray, temp, bounce_count+1);
      
      refres.Scale(reflectiveColor.x(),reflectiveColor.y(),reflectiveColor.z());
      
      answer += refres;
      
      RayTree::AddReflectedSegment(ray, 0, temp.getT());
    }
  


  // =================================
  // ASSIGNMENT:  ADD REFLECTIVE LOGIC
  // =================================


  
  return answer; 
}
